import re
from api.base import *
from common.logger import logger
from common.retcode import RetCode
from common.encrypt import encrypt, b64tohex

class LoginAction(BaseAction):
	def __init__(self, sess_wp: SessionWrapper):
		super(LoginAction, self).__init__(sess_wp)

		self._captcha_handler = None

	def _needcaptcha(self, username):
		"""登录验证码处理函数"""
		url = self._auth_url + "needcaptcha.do"
		post_data = {
			"accountType": "01",
			"userName": "{RSA}" + b64tohex(encrypt(username)),
			"appKey": "cloud"
		}
		r = self._post(url, data=post_data)
		captcha = ""
		if r.text != "0":  # 需要验证码
			if self._captcha_handler:
				img_data = self.refreshCaptcha() 
				captcha = self._captcha_handler(img_data)  # 用户手动识别验证码
			else:
				logger.error("No verification code processing function!")
		return captcha

	def set_captcha_handler(self, captcha_handler):
		"""设置下载验证码处理函数
		:param captcha_handler (img_data) -> str 参数为图片二进制数据,需返回验证码字符
		"""
		self._captcha_handler = captcha_handler

	def refreshCaptcha(self):
		pic_url = self._auth_url + "picCaptcha.do"
		img_data = self._get(pic_url, params={"token": self._sess_wp.get_captcha_token()}).content
		return img_data

	def login(self, username, password):
		"""使用 用户名+密码 登录"""

		# 先清cookies，否则访问 /udb/udb_login.jsp 会返回重定向页面
		self._sess_wp.clear_all_cookies()

		url = self._host_url + "/udb/udb_login.jsp"
		params = {"pageId": 1, "redirectURL": "/main.action"}
		resp = self._get(url, params=params)
		if not resp:
			logger.error("redirect error!")
			return RetCode.NETWORK_ERROR
		# captchaToken = re.search(r"captchaToken' value='(.+?)'", resp.text)
		captchaToken = re.search(r"captchaToken\W*value=\W*(\w*)", resp.text)
		# returnUrl = re.search(r"returnUrl = '(.+?)'", resp.text)
		returnUrl = re.search(r"returnUrl =\W*([^'\"]*)", resp.text)
		# paramId = re.search(r'paramId *=\W*(\w*)', resp.text)
		paramId = re.search(r'paramId =\W*(\w*)', resp.text)
		# lt = re.search(r'lt = "(.+?)"', resp.text)
		lt = re.search(r'lt =\W+(\w*)', resp.text)

		captchaToken = captchaToken.group(1) if captchaToken else ""
		lt = lt.group(1) if lt else ""
		returnUrl = returnUrl.group(1) if returnUrl else ""
		paramId = paramId.group(1) if paramId else ""
		logger.debug(f"Login: {captchaToken}, {lt}, {returnUrl}, {paramId}")
		self._sess_wp.set_header({"lt": lt})
		self._sess_wp.set_captcha_token(captchaToken)

		validateCode = self._needcaptcha(username)
		if validateCode == None:
			return RetCode.CAPTCHA_ERROR
		url = self._auth_url + "loginSubmit.do"
		data = {
			"appKey": "cloud",
			"accountType": '01',
			"userName": "{RSA}" + b64tohex(encrypt(username)),
			"password": "{RSA}" + b64tohex(encrypt(password)),
			"validateCode": validateCode,
			"captchaToken": captchaToken,
			"returnUrl": returnUrl,
			"mailSuffix": "@189.cn",
			"paramId": paramId
		}
		r = self._post(url, data=data)
		r_json = r.json()
		logger.debug(r_json)
		if r_json["result"] == 0:
			self._get(r_json["toUrl"])
			return RetCode.SUCCESS

		return RetCode.FAILED

	def login_by_cookie(self):
		try:
			resp = self._get(self._host_url + "/v2/getUserLevelInfo.action")
			# print(f'login_by_cookie: {resp.text}')
			resp_json = resp.json()
			print(f'login_by_cookie json: {resp_json}')
			if "InvalidSessionKey" not in resp.text and ('data' in resp_json):
				return RetCode.SUCCESS
		except:
			pass
		return RetCode.FAILED